from django.contrib.auth import get_user_model
from channels.generic.websocket import WebsocketConsumer


class CountConsumer(WebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(args, kwargs)
        self.user = None

    def connect(self):
        self.user = self.scope['user']
        self.accept()

    def receive(self, text_data=None, bytes_data=None):
        from chat.models import Message
        sender = get_user_model().objects.filter(username=text_data).first()
        if sender is None:
            messages = Message.objects.filter(receiver_id=self.user.id, read=False)
        else:
            messages = Message.objects.filter(sender_id=sender.id, receiver_id=self.user.id, read=False)
        self.send(text_data=str(messages.count()))

    def disconnect(self, code):
        self.close()
